%% copy riak-0.14.2 source to learn dynamo implementation and erlang
%% author: lichuang
%% copy source from riak_kv_pb_socket.erl
-module(edynamo_pb_socket).
-behaviour(gen_server).

-include_lib("../include/riakclient_pb.hrl").
-include_lib("../include/riakc_pb.hrl").
-export([start_link/0, set_socket/2]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
    terminate/2, code_change/3]).

-type msg() ::  atom() | tuple().

-record(state, {sock,      % protocol buffers socket
	client,    % local client
	req,       % current request (for multi-message requests like list keys)
	req_ctx}). % context to go along with request (partial results, request ids etc)
%% ===================================================================
%% Public API
%% ===================================================================
start_link() ->
  gen_server:start_link(?MODULE, [], []).

set_socket(Pid, Socket) ->
  gen_server:call(Pid, {set_socket, Socket}).

init(_Args) ->
  %edynamo_stat:update(pbc_connect),
  {ok, C} = edynamo_util:local_client(),
  {ok, #state{client = C}}.

handle_call({set_socket, Socket}, _From, State) ->
  inet:setopts(Socket, [{active, once}, {packet, 4}, {header, 1}]),
  {reply, ok, State#state{sock = Socket}};

handle_call(_Args, _From, State) ->
  {reply, ok, State}.

handle_cast(_Msg, State) -> 
  {noreply, State}.

handle_info({tcp_closed, Socket}, State=#state{sock = Socket}) ->
  {stop, normal, State};
handle_info({tcp_error, Socket, _Reason}, State=#state{sock = Socket}) ->
  {stop, normal, State};
handle_info({tcp, _Sock, Data}, State=#state{sock = Socket}) ->
  [MsgCode|MsgData] = Data,
  Msg = riakc_pb:decode(MsgCode, MsgData),
  case process_message(Msg, State) of
    {pause, NewState} ->
      ok;
    NewState ->
      inet:setopts(Socket, [{active, once}])
  end,
  {noreply, NewState};

handle_info(_, State) -> % Ignore any late replies from gen_servers/messages from fsms
  {noreply, State}.

terminate(_Reason, _State) -> 
  %riak_kv_stat:update(pbc_disconnect),
  ok.

code_change(_OldVsn, State, _Extra) ->
  {ok, State}.

%% ===================================================================
%% Message Handling
%% ===================================================================

%% Process an incoming protocol buffers message.  Return either
%% a new #state{} if new incoming messages should be received
%% or {pause, #state{}} if the incoming TCP socket should not be
%% set active again.
%%
%% If 'pause' is returned, it needs to be re-enabled by whatever
%% callbacks are waiting for it.
%%
-spec process_message(msg(), #state{}) -> #state{} | {pause, #state{}}.
process_message(rpbpingreq, State) ->
  send_msg(rpbpingresp, State);

process_message(#rpbgetreq{bucket=B, key=K, r=R0}, 
                #state{client=C} = State) ->
  R = normalize_rw_value(R0),
  case C:get(B, K, default_r(R)) of
    {ok, O} ->
      PbContent = riakc_pb:pbify_rpbcontents(edynamo_object:get_contents(O), []),
      GetResp = #rpbgetresp{content = PbContent,
                            vclock  = pbify_rpbvc(edynamo_object:vclock(O))},
      send_msg(GetResp, State);
    {error, notfound} ->
      send_msg(#rpbgetresp{}, State);
    {error, Reason} ->
      send_error("~p", [Reason], State)
  end;

process_message(#rpbputreq{bucket=B, key=K,vclock=PbVc,content=RpbContent,
    w=W0, dw=DW0, return_body=ReturnBody}, #state{client=C} = State) ->
  O0 = edynamo_object:new(B, K, <<>>),
  error_logger:error_msg("RpbContent:~p~n", [RpbContent]),
  error_logger:error_msg("PbVc:~p~n", [PbVc]),
  O1 = update_rpbcontent(O0, RpbContent),
  O  = update_pbvc(O1, PbVc),
  % erlang_protobuffs encodes as 1/0/undefined
  W  = normalize_rw_value(W0),
  DW = normalize_rw_value(DW0),
  Options = case ReturnBody of 1 -> [returnbody]; _ -> [] end,
  case C:put(O, default_w(W), default_dw(DW), default_timeout(), Options) of
    ok ->
      send_msg(#rpbputresp{}, State);
    {ok, Obj} ->
      PbContents = riakc_pb:pbify_rpbcontents(edynamo_object:get_contents(Obj),[]),
      PutResp = #rpbputresp{content = PbContents,
	vclock  = pbify_rpbvc(edynamo_object:vclock(Obj))},
      send_msg(PutResp, State);
    {error, notfound} ->
      send_msg(#rpbputresp{}, State);
    {error, Reason} ->
      send_error("~p", [Reason], State)
  end.

%% Update edynamo_object with the pbcontent provided
update_rpbcontent(O0, RpbContent) ->
  {MetaData, Value} = riakc_pb:erlify_rpbcontent(RpbContent),
  error_logger:error_msg("meta:~p, value:~p~n", [MetaData, Value]),
  O1 = edynamo_object:update_metadata(O0, MetaData),
  edynamo_object:update_value(O1, Value).

%% Update edynamo_object with vector clock 
update_pbvc(O0, PbVc) ->
  Vclock = erlify_rpbvc(PbVc),
  edynamo_object:set_vclock(O0, Vclock).

%% Send a message to the client
send_msg(Msg, State) ->
  Pkt = riakc_pb:encode(Msg),
  gen_tcp:send(State#state.sock, Pkt),
  State.

%% Send an error to the client
-spec send_error(string(), list(), #state{}) -> #state{}.
send_error(Msg, Fmt, State) ->
  send_error(Msg, Fmt, ?RIAKC_ERR_GENERAL, State).

-spec send_error(string(), list(), non_neg_integer(), #state{}) -> #state{}.
send_error(Msg, Fmt, ErrCode, State) ->
  %% protocol buffers accepts nested lists for binaries so no need to flatten the list
  ErrMsg = io_lib:format(Msg, Fmt),
  send_msg(#rpberrorresp{errmsg=ErrMsg, errcode=ErrCode}, State).

%% Set default values in the options record if none are provided.
%% Erlang protobuffs does not support default, so have to do it here.
default_r(undefined) ->
  default;
default_r(R) ->
  R.

default_w(undefined) ->
  default;
default_w(W) ->
  W.

default_dw(undefined) ->
  default;
default_dw(DW) ->
  DW.

default_rw(undefined) ->
  default;
default_rw(RW) ->
  RW.

default_timeout() ->
  60000.

%% Convert a vector clock to erlang
erlify_rpbvc(undefined) ->
  vclock:fresh();
erlify_rpbvc(<<>>) ->
  vclock:fresh();
erlify_rpbvc(PbVc) ->
  binary_to_term(zlib:unzip(PbVc)).

%% Convert a vector clock to protocol buffers
pbify_rpbvc(Vc) ->
  zlib:zip(term_to_binary(Vc)).

normalize_rw_value(?RIAKC_RW_ONE) -> one;
normalize_rw_value(?RIAKC_RW_QUORUM) -> quorum;
normalize_rw_value(?RIAKC_RW_ALL) -> all;
normalize_rw_value(?RIAKC_RW_DEFAULT) -> default;
normalize_rw_value(V) -> V.
